{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3d9f39e4-2a86-45bd-abac-98117aefa7ac",
   "metadata": {},
   "source": [
    "# Anthropic Tool calling capabilities with Unity Catalog\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "**API Key**\n",
    "To run this tutorial, you will need an Anthropic API key. For testing purposes, you can generate a new account, and use your evaluation test key (no credit card required!).\n",
    "\n",
    "Once you have acquired your key, set it to the environment variable `ANTHROPIC_API_KEY`.\n",
    "\n",
    "Below, we validate that this key is set properly in your environment.\n",
    "\n",
    "**Packages**\n",
    "\n",
    "To interface with both UnityCatalog and Anthropic, you will need to install the following packages:\n",
    "\n",
    "```shell\n",
    "pip install anthropic unitycatalog-anthropic\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "580c0b0f-3ab8-4d5a-9dea-0ee8c516c566",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "assert \"ANTHROPIC_API_KEY\" in os.environ, (\n",
    "    \"Please set the ANTHROPIC_API_KEY environment variable to your Anthropic API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7fee661-4a32-425e-b42c-6f4efcb2eec8",
   "metadata": {},
   "source": [
    "## Configuration and Client setup\n",
    "\n",
    "In order to connect to your Unity Catalog server, you'll need an instance of the `ApiClient` from the `unitycatalog-client` package. \n",
    "\n",
    "> Note: If you don't already have a Catalog and a Schema created, be sure to create them before running this notebook and adjust the `CATALOG` and `SCHEMA` variables below to suit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eba0f1bf-ef95-420e-91cd-392e0a4a509f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from unitycatalog.ai.anthropic.toolkit import UCFunctionToolkit\n",
    "from unitycatalog.ai.anthropic.utils import generate_tool_call_messages\n",
    "from unitycatalog.ai.core.client import UnitycatalogFunctionClient\n",
    "from unitycatalog.client import ApiClient, Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9a93a22f-4d55-428e-8576-ee11be1c009d",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Configuration()\n",
    "config.host = \"http://localhost:8080/api/2.1/unity-catalog\"\n",
    "\n",
    "# The base ApiClient is async\n",
    "api_client = ApiClient(configuration=config)\n",
    "\n",
    "client = UnitycatalogFunctionClient(api_client=api_client)\n",
    "\n",
    "CATALOG = \"AICatalog\"\n",
    "SCHEMA = \"AISchema\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d2ecec7-b8be-485b-9ead-157cb9265aad",
   "metadata": {},
   "source": [
    "## Define a function and register it to Unity Catalog\n",
    "\n",
    "In this next section, we'll be defining a placeholder Python function and creating it within Unity Catalog so that it can be retrieved and used as a tool by Anthropic's Claude model. \n",
    "\n",
    "There are a few things to keep in mind when creating functions for use with the `create_python_function` API:\n",
    "\n",
    "- Ensure that your have properly defined types for all arguments and for the return of the function.\n",
    "- Ensure that you have a Google-style docstring defined that includes descriptions for the function, each argument, and the return of the function. This is critical, as these are used to populate the metadata associated with the function within Unity Catalog, providing contextual data for an LLM to understand when and how to call the tool associated with this function.\n",
    "- If there are packages being called that are not part of core Python, ensure that the import statements are locally scoped (defined within the function body)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "99d7c4fa-c930-4824-aa2e-2ccf5ad73320",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fetch_weather(location: str) -> str:\n",
    "    \"\"\"\n",
    "    Fetches the current weather in celsius for a given location.\n",
    "\n",
    "    Args:\n",
    "        location (str): The location to fetch the weather for.\n",
    "\n",
    "    Returns:\n",
    "        str: The current weather in celsius for the given location.\n",
    "    \"\"\"\n",
    "\n",
    "    return \"88.2 F\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9e2ecc06-78fa-423b-b6f4-86a649afb6cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionInfo(name='fetch_weather', catalog_name='AICatalog', schema_name='AISchema', input_params=FunctionParameterInfos(parameters=[FunctionParameterInfo(name='location', type_text='STRING', type_json='{\"name\": \"location\", \"type\": \"string\", \"nullable\": false, \"metadata\": {\"comment\": \"The location to fetch the weather for.\"}}', type_name=<ColumnTypeName.STRING: 'STRING'>, type_precision=None, type_scale=None, type_interval_type=None, position=0, parameter_mode=None, parameter_type=None, parameter_default=None, comment='The location to fetch the weather for.')]), data_type=<ColumnTypeName.STRING: 'STRING'>, full_data_type='STRING', return_params=None, routine_body='EXTERNAL', routine_definition='return \"88.2 F\"', routine_dependencies=None, parameter_style='S', is_deterministic=True, sql_data_access='NO_SQL', is_null_call=False, security_type='DEFINER', specific_name='fetch_weather', comment='Fetches the current weather in celsius for a given location.', properties='null', full_name='AICatalog.AISchema.fetch_weather', owner=None, created_at=1732579597979, created_by=None, updated_at=1732579597979, updated_by=None, function_id='70623a66-0775-4567-92f6-47b0b1d6b288', external_language='PYTHON')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.create_python_function(func=fetch_weather, catalog=CATALOG, schema=SCHEMA, replace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3dbc6c5-362b-40fb-b561-f9f8c588793b",
   "metadata": {},
   "source": [
    "## Create a Toolkit instance of the function(s)\n",
    "\n",
    "Now that the function has been created within Unity Catalog, we can use the `unitycatalog-anthropic` package to create a toolkit instance that Anthropic will 'understand' as a valid tool in its APIs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3d948b39-fa0b-4338-ba5d-7eba40ce947b",
   "metadata": {},
   "outputs": [],
   "source": [
    "toolkit = UCFunctionToolkit(function_names=[f\"{CATALOG}.{SCHEMA}.fetch_weather\"], client=client)\n",
    "\n",
    "tools = toolkit.tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "81f780f3-02ea-454f-9578-2360c483dc34",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[AnthropicTool(name='AICatalog__AISchema__fetch_weather', description='Fetches the current weather in celsius for a given location.', input_schema={'type': 'object', 'properties': {'location': {'default': None, 'description': 'The location to fetch the weather for.', 'title': 'Location', 'type': 'string'}}, 'required': []})]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5f27777-e136-4cac-bf6b-1a9bbc89ca95",
   "metadata": {},
   "source": [
    "## Call Claude with tools\n",
    "\n",
    "With our toolkit instance ready, we can now pass in our tools to a call to Anthropic, giving Claude the ability to request for 'answers' from tools that have been provided with the request."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "617a5de6-a108-425c-a710-da4286efb075",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Message(id='msg_012UX8sMYq55q1mR3HzbaWDL', content=[TextBlock(text=\"To answer your question about the weather in Nome, AK and Death Valley, CA, I'll need to use the weather fetching tool twice, once for each location. Let me do that for you now.\", type='text'), ToolUseBlock(id='toolu_01MRhL9hHEcYsSBoM1hXXeYb', input={'location': 'Nome, AK'}, name='AICatalog__AISchema__fetch_weather', type='tool_use'), ToolUseBlock(id='toolu_01NwcjWquf1bJvzjwxyi43H2', input={'location': 'Death Valley, CA'}, name='AICatalog__AISchema__fetch_weather', type='tool_use')], model='claude-3-5-sonnet-20240620', role='assistant', stop_reason='tool_use', stop_sequence=None, type='message', usage=Usage(input_tokens=406, output_tokens=155))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Interface with Anthropic via their SDK\n",
    "import anthropic\n",
    "\n",
    "anthropic_client = anthropic.Anthropic()\n",
    "\n",
    "multi = \"What's the weather in Nome, AK and in Death Valley, CA?\"\n",
    "\n",
    "question = [{\"role\": \"user\", \"content\": multi}]\n",
    "\n",
    "response = anthropic_client.messages.create(\n",
    "    model=\"claude-3-5-sonnet-20240620\",\n",
    "    max_tokens=1024,\n",
    "    tools=tools,\n",
    "    messages=question,\n",
    ")\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d91e129-3579-4cd5-9ad7-b40cc44fc407",
   "metadata": {},
   "source": [
    "## Use the helper utility to call the tools and continue the conversation\n",
    "\n",
    "With the response from Claude being a request to execute specific tools with parameters, we can utilize the `generate_tool_call_messages` helper utility from `unitycatalog.ai.anthropic.utils`. This function will process the response from Anthropic, parse the data, determine whether tool(s) need to be called, call the tool(s), and formulate the entire response structure for submitting the follow-on response call back to Anthropic.\n",
    "\n",
    "> Note: If not using this helper functionality, keep in mind that Anthropic will submit multiple tool request calls if it is deemed as necessary. If complex multi-turn tool calls are needed (i.e. a response from `Tool A` is needed to call `Tool B`, additional `ToolUseBlock` responses will be generated on a subsequent call), you will need to handle this yourself. The utility function, on the other hand, is capable of handling this logic and will continue to process tool calls until the response structure contains no such elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9a080f0e-74ad-419b-9c85-39a5d4a74f35",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Message(id='msg_01Vjv77qn7LxVu4EwozDSjQ6', content=[TextBlock(text=\"Thank you for waiting. I've fetched the current weather information for both locations. Here's what I found:\\n\\n1. Nome, AK: The current temperature is 88.2°F (31.2°C).\\n2. Death Valley, CA: The current temperature is also 88.2°F (31.2°C).\\n\\nInterestingly, both locations are showing the same temperature right now, which is quite unusual given their very different geographical locations and typical climate patterns. Nome, Alaska is typically much cooler than Death Valley, California, especially during most of the year.\\n\\nDeath Valley is known for its extremely hot temperatures, especially during summer months, so this temperature is not unusual for that location. However, for Nome, Alaska, this is an exceptionally high temperature, as it's located near the Arctic Circle and generally experiences much cooler weather.\\n\\nIf you'd like more detailed weather information or have any questions about this unusual weather pattern, please feel free to ask.\", type='text')], model='claude-3-5-sonnet-20240620', role='assistant', stop_reason='end_turn', stop_sequence=None, type='message', usage=Usage(input_tokens=649, output_tokens=212))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Call the Unity Catalog function and construct the required formatted response history for a subsequent call to Anthropic\n",
    "tool_messages = generate_tool_call_messages(\n",
    "    response=response, client=client, conversation_history=question\n",
    ")\n",
    "\n",
    "# Call the Anthropic client with the parsed tool response from executing the Unity Catalog function\n",
    "tool_response = anthropic_client.messages.create(\n",
    "    model=\"claude-3-5-sonnet-20240620\",\n",
    "    max_tokens=1024,\n",
    "    tools=tools,\n",
    "    messages=tool_messages,\n",
    ")\n",
    "\n",
    "tool_response"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
